import torch


def compute_periods(
    period_min: torch.Tensor | float,
    period_max: torch.Tensor | float,
    periods_n: int,
    dtype: torch.dtype | None = None,
    device: torch.device | None = None,
) -> torch.Tensor:
    period_min = torch.asarray(period_min, dtype=dtype, device=device)
    period_max = torch.asarray(period_max, dtype=dtype, device=device)
    period_min_log = torch.log(period_min)
    period_max_log = torch.log(period_max)
    indices = torch.arange(periods_n, dtype=dtype, device=device)
    periods_log = period_min_log + (period_max_log - period_min_log) * indices / (periods_n - 1)
    periods = torch.exp(periods_log)
    return periods


def compute_periods_cs(*,
    pos: torch.Tensor,
    size: int,
    period_min: torch.Tensor | float,
    period_max: torch.Tensor | float,
    dtype: torch.dtype | None = None,
    device: torch.device | None = None,
) -> torch.Tensor:
    pos = torch.asarray(pos, dtype=torch.float64, device=device)
    pos = torch.unsqueeze(pos, dim=-1)
    periods = compute_periods(period_min, period_max, size // 2, dtype=torch.float64, device=device)
    periods_c = torch.cos(pos / periods * (2 * torch.pi))
    periods_s = torch.sin(pos / periods * (2 * torch.pi))
    periods_cs = torch.stack((periods_c, periods_s), dim=-1)
    periods_cs = torch.reshape(periods_cs, (*periods_cs.shape[:-2], -1))
    periods_cs = periods_cs.to(dtype=dtype)
    return periods_cs


def apply_rotary_emb(
    x: torch.Tensor,
    periods_cs: torch.Tensor,
) -> torch.Tensor:
    x = torch.reshape(x, (*x.shape[:-1], -1, 2))
    x_c, x_s = x.unbind(dim=-1)
    periods_cs = torch.reshape(periods_cs, (*periods_cs.shape[:-1], -1, 2))
    periods_c, periods_s = periods_cs.unbind(dim=-1)
    x_out_c = x_c * periods_c - x_s * periods_s
    x_out_s = x_c * periods_s + x_s * periods_c
    x_out = torch.stack((x_out_c, x_out_s), dim=-1)
    return x_out.reshape((*x.shape[:-2], -1))
